import user  # look for .pythonrc.py for user init
import math
import ROOT
import PyCintex
import AthenaROOTAccess.transientTree

def getBackend():
   '''returns Grid backend: 'Panda' or 'LCG' or 'None'
   '''
   import commands
   ##### are we running on the Grid?
   ### Ganga's LCG backend produces a file input.txt
   ### Ganga's Panda backend has pre_*.py and post_*.py
   backend = 'None'
   cmd = 'ls input.txt'
   status, output = commands.getstatusoutput(cmd)
   if output.find('No such file') == -1:
      backend = 'LCG'
   else:
      cmd = 'ls pre_*.py'
      status, output = commands.getstatusoutput(cmd)
      if output.find('No such file') == -1:
         backend = 'Panda'

   return backend
  

backend = getBackend()

CollectionTree = \
            ROOT.AthenaROOTAccess.TChainROOTAccess('CollectionTree')

# if running on the LCG, the list of input files is read from the file 
# input.txt which is placed in the run directory by Ganga before the
# program is called
if backend == 'LCG':
   for myfile in open("input.txt","r"):
      inputfile = myfile.strip()
      # do not use log files
      if inputfile.find('log.tgz') == -1 and inputfile.find('root') != -1:
         print 'using: ' + inputfile
         CollectionTree.Add( inputfile )
            
# for running on the PANDA grid, use all *.root* files in run directory as
# input    
elif backend == 'Panda':
   import glob
   l_input = glob.glob('*.root*')
   l_input.sort()
   for inputfile in l_input:
      print 'using: ' + inputfile
      CollectionTree.Add( inputfile )
      
# for local running specify location of input files   
else:   
   CollectionTree.Add('/tmp/narayan/mc11_7TeV.109300.AlpgenJimmyZeebbNp0_nofilter.evgen.EVNT.e835_tid454677_00/*root*')

tt = AthenaROOTAccess.transientTree.makeTree(CollectionTree)
#############################
#############################
#############################
def toiter (beg, end):
       while beg != end:
           yield beg.__deref__()
           beg.__preinc__()
       return
def DeltaR(p,q):
    from math import sqrt
    return sqrt((p.momentum().eta() - q.momentum().eta())**2 + (p.momentum().phi() - q.momentum().phi())**2)

def IsBHadron(p):
    status = p.status()
    mpdg = abs(p.pdg_id())
    if(status==196 or status == 197 or status ==198):
        if( ( 500 < mpdg and mpdg < 599 )    or \
                ( 10500 < mpdg and mpdg < 10599 )  or \
                (  5000 < mpdg and mpdg < 5999  )  or \
                ( 20500 < mpdg and mpdg < 20599 ) ):
                return p

def IsCHadron(p):
    status = p.status()
    mpdg = abs(p.pdg_id())
    if(status == 196 or status == 197 or status ==198): 
        if( ( 400  < mpdg and mpdg < 499 )   or \
                ( 10400 < mpdg and mpdg < 10499 ) or \
                (  4000 < mpdg and mpdg < 4999  ) or \
                ( 20400 < mpdg and mpdg < 20499 ) ):
            return p



def IsMEParton(p):
    hasbchadronparent   = False
    hasbcquarkdaughter  = False
    hasbquarkparent     = False
    hasmpiparent        = False
    hastopparent        = False
    haswparent          = False

    isMEParton = False
    apdg    = abs(p.pdg_id())
    if (apdg == 5 or apdg ==4):
        prodVtx = p.production_vertex()
        decayVtx= p.end_vertex()
        if(prodVtx):
            ''' Loop over the parents of the current particle'''
            for parent in toiter(prodVtx.particles_begin(0),prodVtx.particles_end(0)): 
                '''check if there is no b/c hadron parent'''
                if((abs(parent.pdg_id())%10000)/1000 ==  apdg  or (abs(parent.pdg_id())%1000)/100 == apdg):
                    hasbchadronparent = True
                    break
                ''' reject c-quarks from a b-quark /hadron decay'''
                if ( apdg == 4 and ( abs(parent.pdg_id()) == 5 or  (abs(parent.pdg_id())%10000)/1000 == 5 or (abs(parent.pdg_id())%1000)/100 == 5 ) ):
                        hasbquarkparent = True
                        break
                ''' reject particle with MPI parent '''
                if(abs(parent.pdg_id()) ==0 or parent.status() == 120): #Herwig specific status codes 
                    hasmpiparent= True

                if(parent.pdg_id()==6):
                    hastopparent = True

                if(parent.pdg_id()==24):
                    haswparent = True

        if(not hasbchadronparent):
            if( not hasmpiparent):
                if( not hasbquarkparent):
                    if(p.status() ==123 or p.status()==124):
                        isMEParton = True

    return isMEParton



hists = {}

hists["h1_nEvents"]     = ROOT.TH1F("nEvents","",4,0,4)

#Hard HF events (ME)
hists["h1_nHardbEvents"]= ROOT.TH1F("nHardbEvents","",4,0,4)
hists["h1_nHardcEvents"]= ROOT.TH1F("nHardcEvents","",4,0,4)

#HF hadrons
hists["h1_nBHadron"]        = ROOT.TH1F("nBHadron","",4,0,4)
hists["h1_nBHadronSemilept"]= ROOT.TH1F("nBHadronSemilept","",4,0,4)
hists["h1_nHardB"]          = ROOT.TH1F("nHardB","",4,0,4)
hists["h1_nHardBSemilept"]  = ROOT.TH1F("nHardBSemilept","",4,0,4)

hists["h1_nCHadron"]        = ROOT.TH1F("nCHadron","",4,0,4)
hists["h1_nCHadronSemilept"]= ROOT.TH1F("nCHadronSemilept","",4,0,4)
hists["h1_nHardC"]          = ROOT.TH1F("nHardC","",4,0,4)
hists["h1_nHardCSemilept"]  = ROOT.TH1F("nHardCSemilept","",4,0,4)

# n Evetns Q->Hadron->El
hists["h1_nBHElEvents"] = ROOT.TH1F("nBHElEvents","",4,0,4)
hists["h1_nCHElEvents"] = ROOT.TH1F("nCHElEvents","",4,0,4)

#HF Hadron Decay Electron
hists["h1_BElPt"]       = ROOT.TH1F("BElPt","",500,0,20)
hists["h1_HardBElPt"]   = ROOT.TH1F("HardBElPt","",500,0,20)
hists["h1_CElPt"]       = ROOT.TH1F("CElPt","",500,0,20)
hists["h1_HardCElPt"]   = ROOT.TH1F("HardCElPt","",500,0,20)

#N Hard El vs n Soft El
hists["h2_ElHardVsSoft"]= ROOT.TH2F("ElHardVsSoft",";# Hard el; #Soft el",20,0,20,20,0,20)
hists["h2_ElHardVsSoftTrue"]= ROOT.TH2F("ElHardVsSoftTrue",";# Hard el; #Soft el",20,0,20,20,0,20)

#2D histogram Quark pt Hadron delta r
hists["h2_BMatching"]   = ROOT.TH2F("BMatching",";p_{T}[GeV];#Delta R",500,0,500,100,0,4)
hists["h2_CMatching"]   = ROOT.TH2F("CMatching",";p_{T}[GeV];#Delta R",500,0,500,100,0,4)

#Matched Delta r
hists["h1_MatchedBdr"]  = ROOT.TH1F("MatchedBdr","",100,0,4)
hists["h1_MatchedCdr"]  = ROOT.TH1F("MatchedCdr","",100,0,4)

#Number of HF Q-hadron match per event
hists["h1_nBMatches"]   = ROOT.TH1F("nBMatches",";# Matches",10,0,10)
hists["h1_nCMatches"]   = ROOT.TH1F("nCMatches",";# Matches",10,0,10)

#Delta R of the quarks as a function of Quark Pt
hists["h3_MEbQuarkDr"]  = ROOT.TH3F("MEbQuarkDeltaR","",100,0,4,500,0,500,500,0,500)
hists["h3_MEcQuarkDr"]  = ROOT.TH3F("MEcQuarkDeltaR","",100,0,4,500,0,500,500,0,500)

#Electron kinematics
hists["h1_ElPt"]        = ROOT.TH1F("ElPt",";[GeV]",500,0,500)
hists["h1_ElEta"]       = ROOT.TH1F("ElEta","",200,-5,5)
hists["h1_ElPhi"]       = ROOT.TH1F("ElPhi","",200,-5,5)

[h.Sumw2() for name,h in hists.iteritems()]

for iEvt in range(tt.GetEntriesFast()):
#for iEvt in range(500):
    print "Event"
    if(not iEvt %100):
        print "Event", iEvt
    tt.GetEntry(iEvt)
    ev = tt.GEN_EVENT[0]
    #ev = tt.GEN_AOD[0]
    #mcc = tt.SpclMC
    #ev = mcc.genEvent()
    
    vitr = ev.vertices_begin()
    pVtx = ROOT.TVector3(vitr.position().x(),vitr.position().y(), vitr.position().z())
    
    nElectrons = 0
    nHardElectrons = 0
    nHardZElectrons = 0
    nSoftElectrons = 0
    nSoftBElectrons = 0

    nBMatches = 0
    nCMatches = 0
    
    HardbPartonList = []
    HardcPartonList =[]
 
    MEBHadronList   = []
    PSBHadronList   = []
    BHadronList     = []
    BHadronDaughterEl= []

 
    MECHadronList   = []
    PSCHadronList   = []
    CHadronList     = []
    CHadronDaughterEl=[]

    isbHEl = False
    iscHEl = False

    hists["h1_nEvents"].Fill(1)
    for p in toiter(ev.particles_begin(), ev.particles_end()):
        '''Make a list of ME HF partons'''
        if(IsMEParton(p)):
            if(abs(p.pdg_id())==5):
                HardbPartonList +=[p]
            if(abs(p.pdg_id())==4):
                HardcPartonList +=[p]

    for p in toiter(ev.particles_begin(), ev.particles_end()):
        if(IsBHadron(p)):
            hists["h1_nBHadron"].Fill(1)
            BHadronList += [p]

            dr = -1
            isMEHadron = False
            for MEparton in HardbPartonList:
                dr = DeltaR(p,MEparton)
                if(dr < 0.5):
                    isMEHadron = True
                    break 

            if(isMEHadron):
                MEBHadronList += [p]
                hists["h1_nHardB"].Fill(1)
                hists["h1_MatchedBdr"].Fill(dr)
                nBMatches +=1
            else:
                PSBHadronList += [p]

            if(p.end_vertex()):
                daughterIsEl = False
                for daughter in toiter(p.end_vertex().particles_begin(1),p.end_vertex().particles_end(1)):
                    if(abs(daughter.pdg_id())==11):
                        daughterIsEl = True
                        BHadronDaughterEl += [daughter]
                        hists["h1_BElPt"].Fill(daughter.momentum().perp()/1000)
                        if(isMEHadron):
                            hists["h1_HardBElPt"].Fill(daughter.momentum().perp()/100)

                if(daughterIsEl):
                    hists["h1_nBHadronSemilept"].Fill(1)
                    if(isMEHadron):
                        hists["h1_nHardBSemilept"].Fill(1)
                        isbHEl = True


        elif(IsCHadron(p)):
            hists["h1_nCHadron"].Fill(1)
            CHadronList += [p]

            dr = -1
            isMEHadron = False
            for MEparton in HardcPartonList:
                dr = DeltaR(p,MEparton)
                if(dr < 0.3):
                    isMEHadron = True
                    break

            if(isMEHadron):
                MECHadronList += [p]
                hists["h1_nHardC"].Fill(1)
                hists["h1_MatchedCdr"].Fill(dr)
                nCMatches +=1
            else:
                PSCHadronList += [p]

            if(p.end_vertex()):
                daughterIsEl = False
                for daughter in toiter(p.end_vertex().particles_begin(1),p.end_vertex().particles_end(1)):
                    if(abs(daughter.pdg_id())==11 ):
                        daughterIsEl = True
                        CHadronDaughterEl +=[p]
                        hists["h1_CElPt"].Fill(daughter.momentum().perp()/1000)
                        if(isMEHadron):
                            hists["h1_HardCElPt"].Fill(daughter.momentum().perp()/1000)
                
                if(daughterIsEl):
                    hists["h1_nCHadronSemilept"].Fill(1)
                    if(isMEHadron):
                        hists["h1_nHardCSemilept"].Fill(1)
                        iscHEl = True

    for p in toiter(ev.particles_begin(), ev.particles_end()):
        if(abs(p.pdg_id()) ==11):
            nElectrons +=1
            hists["h1_ElPt"].Fill(p.momentum().perp()/1000)
            hists["h1_ElEta"].Fill(p.momentum().eta())
            hists["h1_ElPhi"].Fill(p.momentum().phi())
            el_pt = p.momentum().perp()/1000
            if(p.production_vertex()):
                elVtx = ROOT.TVector3(\
                        p.production_vertex().position().x(),p.production_vertex().position().y(), p.production_vertex().position().z())
                elDeltaR = elVtx.DeltaR(pVtx)

                if(el_pt>15 and elDeltaR <0.01):
                    if(p.status()==1):
                        nHardElectrons +=1

                    ZElectron =False
                    for elP in toiter(p.production_vertex().particles_begin(0),p.production_vertex().particles_end(0)):
                        if(elP.pdg_id()==23 and (elP.status() == 155 )): #Decaying Z boson
                            ZElectron =True
                            print p.status()
                    if(ZElectron):
                        nHardZElectrons +=1

                if(el_pt > 0.1 and el_pt <= 20):
                    nSoftElectrons += 1
                    BElectron = False
                    for elP in toiter(p.production_vertex().particles_begin(0),p.production_vertex().particles_end(0)):
                        if((elP in MEBHadronList) or (elP in MECHadronList) ):
                            print "SoftElectron"
                            BElectron = True
                    
                    if(BElectron):
                        nSoftBElectrons +=1

    hists["h2_ElHardVsSoft"].Fill(nHardElectrons,nSoftElectrons)
    hists["h2_ElHardVsSoftTrue"].Fill(nHardZElectrons,nSoftBElectrons)


    '''Matching'''
    [hists["h2_BMatching"].Fill(x.momentum().perp()/1000,DeltaR(x,y)) for x in HardbPartonList for y in BHadronList]
    [hists["h2_CMatching"].Fill(x.momentum().perp()/1000,DeltaR(x,y)) for x in HardcPartonList for y in CHadronDaughterEl]

    hists["h1_nBMatches"].Fill(nBMatches)
    hists["h1_nCMatches"].Fill(nCMatches)

    '''ME quark Delta R'''
    [hists["h3_MEbQuarkDr"].Fill(DeltaR(x,y),x.momentum().perp()/1000,y.momentum().perp()/1000) \
            for x in HardbPartonList for y in HardbPartonList if x.pdg_id() != y.pdg_id()]
    [hists["h3_MEcQuarkDr"].Fill(DeltaR(x,y),x.momentum().perp()/1000,y.momentum().perp()/1000) \
            for x in HardcPartonList for y in HardcPartonList if x.pdg_id() != y.pdg_id()]

    if(len(HardbPartonList)):
        hists["h1_nHardbEvents"].Fill(1)
    if(len(HardcPartonList)):
        hists["h1_nHardcEvents"].Fill(1)

    if(isbHEl):
        hists["h1_nBHElEvents"].Fill(1)
    if(iscHEl):
        hists["h1_nCHElEvents"].Fill(1)

f = ROOT.TFile("output.root","RECREATE")
[h.Write() for name, h in hists.iteritems()]
f.Close()
